Walkthrough

This PR removes CUDA-event-based synchronization (_prepare_inputs_event) from PyTorchModelEngine, including its initialization, usage in the forward path, and the wait_for_input_copy() method. It also removes an unrelated test waiver entry from the waives list.

Changes

Prepare-Inputs Synchronization Removal

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested reviewers: StanleySun639, LarryXFly, crazydemo

For the target DS-V3-Lite / V1 KV-cache path, yes: the original behavior is still correct because the race that _prepare_inputs_event papered over was later fixed at the actual mutable host-buffer source by #15546. But I would not claim this removal is generally safe for every V2/DSv4 path unless #16088 or an equivalent V2 staging fix is also present.

Why _prepare_inputs_event was needed

#14128 added _prepare_inputs_event to serialize the overlap scheduler against _prepare_inputs’ async host→device work:

• _prepare_inputs(...) builds attention/page-table inputs and enqueues non-blocking H2D copies.
• In the overlap scheduler, the CPU can start scheduling / updating the next iteration while the previous iteration’s H2D copies are still queued on execution_stream.
• Some H2D sources were persistent pinned host buffers, especially KV block/page-offset tables.
• The next iteration could overwrite those host buffers before the previous H2D copy actually read them, so the GPU could consume block offsets for the wrong batch.

That is why #14128 recorded an event immediately after _prepare_inputs and made the next executor iteration call wait_for_input_copy() before scheduler/resource mutations.

What made it unnecessary for the DS-V3-Lite target

The later fix is #15546: #15546

That PR moved the correctness boundary from a coarse global synchronization to safe per-call host staging in KVCacheManager.copy_batch_block_offsets.

In the current base, KVCacheManager.copy_batch_block_offsets no longer copies directly from the persistent host_kv_cache_block_offsets buffer. It first snapshots the rows needed for this batch into a fresh pinned buffer, then launches the non-blocking H2D from that private snapshot:

• resource_manager.py explains the exact race:
  • async H2D reads source at execution time
  • overlap scheduler can run one iteration ahead
  • persistent host buffer could be clobbered by the next iteration
• Then it stages through _stage_block_offsets_for_copy(...) before the H2D copy.

Relevant code/comments:

# The H2D copies below are asynchronous, so their source is read at copy
# execution time, not enqueue time. With the overlap scheduler the CPU
# runs an iteration ahead, so copying straight from the persistent buffer
# let the next iteration's in-place refill clobber the source of this
# iteration's still-pending H2D -> the attention kernel indexed another
# batch's blocks ...
host_block_offsets = self._stage_block_offsets_for_copy(num_seqs)
for pool_idx in range(self.num_pools):
    dst_tensor[pool_idx, :num_seqs].copy_(host_block_offsets[pool_idx],
                                          non_blocking=True)

That is the key distinction: after #15546, correctness no longer depends on preventing the next iteration from mutating the persistent host table. The in-flight copy reads a private snapshot, so it remains valid even if the next iteration reuses or rewrites the persistent table.

So for the DS-V3-Lite PR #15849 is targeting, the old event is redundant and harmful: it only adds a full CPU-side synchronization point in the overlap path, and that can interact badly with CUDA graph / overlap scheduler progress.

Why this is more than “it won’t crash”

The behavior remains correct because the ordering requirement changed:

• Before [https://nvbugs/6293536][fix] Stage KV block offsets through a fresh host buffer #15546: correctness required “do not mutate persistent host block-offset buffer until previous H2D has consumed it.”
• After [https://nvbugs/6293536][fix] Stage KV block offsets through a fresh host buffer #15546: the previous H2D no longer consumes the persistent buffer; it consumes a per-call snapshot. Therefore the next iteration may safely mutate the persistent buffer.

That preserves the semantic requirement — “the GPU sees the block offsets for the batch being executed” — without requiring a global event synchronize.

Important caveat: V2 / DSv4 is a separate case

The same broad race class exists in V2/DSv4 through a different mechanism. The current DSv4 manager still has paths that copy from persistent host staging into device staging asynchronously, for example:

self._device_kv_cache_block_offsets_input.copy_(
    self.host_kv_cache_block_offsets,
    non_blocking=True,
)

And #16088 explicitly calls out that V2 has the same class of race and needs a V2-specific staging/retirement fix:
#16088

So the strongest accurate statement is:

• Safe for the PR [https://nvbugs/6402058][fix] [None][fix] Remove stale prepare_inputs event sync to unhang DS-V3 EP4 MTP #15849 DS-V3-Lite/V1 target, because [https://nvbugs/6293536][fix] Stage KV block offsets through a fresh host buffer #15546 already fixed the underlying mutable-host-source race.
• Not proven as a general removal for all V2/DSv4 paths unless [https://nvbugs/6293536][fix] Stage v2 KV block offsets through fresh host buffers #16088 or equivalent V2 staging is also in the base.

Suggested PR explanation

I’d phrase the PR rationale like this:

_prepare_inputs_event was originally added as a coarse fence to prevent overlap-scheduler iterations from mutating persistent host input/page-table buffers before previous non-blocking H2D copies consumed them. For the DS-V3-Lite/V1 KV-cache path targeted here, that correctness requirement is now handled locally by KVCacheManager.copy_batch_block_offsets, which snapshots the persistent host block-offset table into a fresh pinned per-call buffer before launching the async H2D copy (#15546 / nvbug 6293536). Therefore the next iteration can safely mutate the persistent table without affecting the in-flight copy, and the global per-forward event synchronize is redundant. Removing it preserves correctness while avoiding an unnecessary overlap-scheduler synchronization that contributes to the nvbug/6402058 hang. Note that V2/DSv4 uses a different staging path and should rely on its own fix (#16088 or equivalent), not on this coarse model-engine event.